import dmPython
import os
import argparse
import logging
import re

def execute_sql_file(host, port, username, password, database, sql_file_path):
    """
    执行达梦SQL文件到指定数据库
    :param host: 数据库主机地址
    :param port: 数据库端口
    :param username: 用户名
    :param password: 密码
    :param database: 数据库名称
    :param sql_file_path: SQL文件路径
    """
    try:
        # 连接到达梦数据库
        conn = dmPython.connect(
            user=username,
            password=password,
            server=host,
            port=port
        )
        if database is not None:
            conn = dmPython.connect(
                user=username,
                password=password,
                server=host,
                port=port,
                database=database
            )
        cursor = conn.cursor()
        
        logging.info(f"成功连接到数据库: {host}:{port}/{database}")
        
        # 读取SQL文件内容
        with open(sql_file_path, 'r', encoding='utf-8') as f:
            sql_script = f.read()
        
        # 分割SQL语句（考虑存储过程等复杂语句）
        sql_commands = split_sql_commands(sql_script)
        
        # 执行每条SQL语句
        success_count = 0
        for i, command in enumerate(sql_commands, 1):
            if command.strip():
                try:
                    cursor.execute(command)
                    conn.commit()
                    success_count += 1
                    logging.info(f"执行成功 [{i}/{len(sql_commands)}]: {command[:50]}...")
                except dmPython.Error as e:
                    logging.error(f"执行失败 [{i}/{len(sql_commands)}]: {command[:50]}...")
                    logging.error(f"错误信息: {e}")
                    conn.rollback()  # 回滚当前事务
        
        logging.info(f"SQL文件执行完成! 成功: {success_count}, 失败: {len(sql_commands)-success_count}")
        
    except Exception as e:
        logging.error(f"数据库操作异常: {e}")
    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals():
            conn.close()

def split_sql_commands(sql_script):
    """
    智能分割SQL脚本中的多条命令
    :param sql_script: SQL脚本内容
    :return: 分割后的SQL命令列表
    """
    # 移除注释
    sql_script = re.sub(r'--.*?$', '', sql_script, flags=re.MULTILINE)
    sql_script = re.sub(r'/\*.*?\*/', '', sql_script, flags=re.DOTALL)
    
    # 分割SQL命令（考虑存储过程等复杂结构）
    commands = []
    current = ""
    delimiter = ';'
    in_string = False
    char_escape = False
    prev_char = None
    
    for char in sql_script:
        if char_escape:
            current += char
            char_escape = False
        elif char == '\\':
            current += char
            char_escape = True
        elif char in ('"', "'"):
            if in_string == char and prev_char != '\\':
                in_string = False
            elif not in_string:
                in_string = char
            current += char
        elif not in_string and char == delimiter:
            commands.append(current.strip())
            current = ""
        else:
            current += char
        prev_char = char
    
    if current.strip():
        commands.append(current.strip())
    
    return commands

def main():
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    
    # 命令行参数解析
    parser = argparse.ArgumentParser(description='达梦SQL文件执行工具')
    parser.add_argument('--host', required=True, help='数据库主机地址')
    parser.add_argument('--port', type=int, default=5236, help='数据库端口 (默认: 5236)')
    parser.add_argument('--user', required=True, help='数据库用户名')
    parser.add_argument('--password', required=True, help='数据库密码')
    parser.add_argument('--database', required=False, help='数据库名称')
    parser.add_argument('--file', required=True, help='SQL文件路径')
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.file):
        logging.error(f"SQL文件不存在: {args.file}")
        return
    
    # 执行SQL文件
    execute_sql_file(
        host=args.host,
        port=args.port,
        username=args.user,
        password=args.password,
        database=args.database,
        sql_file_path=args.file
    )

if __name__ == "__main__":
    main()
